import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from matplotlib.patches import Circle, Rectangle
import math
import time
import random

# 设置递归深度限制（如果需要）
import sys

sys.setrecursionlimit(10000)


class Drone:
    def __init__(self, id, initial_position, communication_range=50, sensor_range=30):
        self.id = id
        self.position = np.array(initial_position, dtype=float)
        self.velocity = np.array([0.0, 0.0], dtype=float)  # 确保是浮点数
        self.max_speed = 3.0
        self.max_acceleration = 1.0
        self.communication_range = communication_range
        self.sensor_range = sensor_range
        self.path = []
        self.waypoints = []
        self.current_waypoint_index = 0
        self.neighbors = []
        self.state = "patrolling"  # 状态: patrolling, scanning, returning, avoiding
        self.battery_level = 100  # 电池电量百分比
        self.battery_drain_rate = 0.05  # 每秒消耗的电量百分比
        self.scan_time = 0
        self.scan_duration = 5  # 扫描持续时间（秒）
        self.shared_info = {}  # 存储从其他无人机共享的信息
        self.obstacles = []  # 存储已知障碍物
        self.targets = []  # 存储发现的目标

    def set_waypoints(self, waypoints):
        """设置巡逻路径点"""
        self.waypoints = waypoints
        self.current_waypoint_index = 0

    def update(self, drones, obstacles, dt):
        """更新无人机状态"""
        # 消耗电池
        self.battery_level -= self.battery_drain_rate * dt
        if self.battery_level <= 20 and self.state != "returning":
            self.state = "returning"

        # 检测附近无人机
        self.detect_neighbors(drones)

        # 检测障碍物
        self.detect_obstacles(obstacles)

        # 共享信息
        self.share_information(drones)

        # 根据状态执行不同行为
        if self.state == "patrolling":
            self.patrol(dt)
        elif self.state == "scanning":
            self.scan_area(dt)
        elif self.state == "returning":
            self.return_to_base(dt)
        elif self.state == "avoiding":
            self.avoid_obstacles(dt)

        # 记录路径（限制路径长度以避免内存问题）
        self.path.append(self.position.copy())
        if len(self.path) > 1000:  # 限制路径长度
            self.path = self.path[-500:]  # 保留最后500个点

    def detect_neighbors(self, drones):
        """检测通信范围内的其他无人机"""
        self.neighbors = []
        for drone in drones:
            if drone.id != self.id:
                distance = np.linalg.norm(self.position - drone.position)
                if distance <= self.communication_range:
                    self.neighbors.append(drone)

    def detect_obstacles(self, obstacles):
        """检测传感器范围内的障碍物"""
        detected_obstacles = []
        for obs in obstacles:
            distance = np.linalg.norm(self.position - obs.position)
            if distance <= self.sensor_range:
                detected_obstacles.append(obs)
                # 添加到已知障碍物列表（如果尚未存在）
                if obs not in self.obstacles:
                    self.obstacles.append(obs)

        # 如果有障碍物在避障范围内，切换到避障状态
        for obs in detected_obstacles:
            if np.linalg.norm(self.position - obs.position) < obs.radius + 15:
                self.state = "avoiding"
                self.avoid_target = obs
                break

    def share_information(self, drones):
        """与邻近无人机共享信息"""
        for neighbor in self.neighbors:
            # 共享障碍物信息
            for obs in self.obstacles:
                if obs not in neighbor.obstacles:
                    neighbor.obstacles.append(obs)

            # 共享目标信息
            for target in self.targets:
                if target not in neighbor.targets:
                    neighbor.targets.append(target)

    def patrol(self, dt):
        """巡逻行为 - 使用改进的路径跟踪算法"""
        if not self.waypoints:
            return

        target = self.waypoints[self.current_waypoint_index]
        direction = target - self.position
        distance = np.linalg.norm(direction)

        if distance < 2.0:  # 到达航点
            self.current_waypoint_index = (self.current_waypoint_index + 1) % len(self.waypoints)
            target = self.waypoints[self.current_waypoint_index]
            direction = target - self.position
            distance = np.linalg.norm(direction)

        # 标准化方向向量
        if distance > 0:
            direction = direction / distance

        # 计算期望速度
        desired_velocity = direction * self.max_speed

        # 应用群体行为（分离、对齐、聚集）
        if self.neighbors:
            separation = self.calculate_separation()
            alignment = self.calculate_alignment()
            cohesion = self.calculate_cohesion()

            # 组合行为（权重可调整）
            desired_velocity = desired_velocity + separation * 1.5 + alignment * 1.0 + cohesion * 1.0

            # 限制速度
            speed = np.linalg.norm(desired_velocity)
            if speed > self.max_speed:
                desired_velocity = desired_velocity / speed * self.max_speed

        # 计算转向力
        steering = desired_velocity - self.velocity
        steering_norm = np.linalg.norm(steering)

        if steering_norm > self.max_acceleration:
            steering = steering / steering_norm * self.max_acceleration

        # 更新速度和位置 - 确保使用浮点数运算
        self.velocity = self.velocity.astype(float) + steering.astype(float) * dt
        self.position = self.position.astype(float) + self.velocity.astype(float) * dt

        # 随机决定是否扫描区域
        if np.random.random() < 0.003:  # 0.3% 的概率
            self.state = "scanning"
            self.scan_time = 0

    def calculate_separation(self):
        """计算分离力（避免与邻居太近）"""
        separation = np.zeros(2, dtype=float)  # 确保是浮点数
        for neighbor in self.neighbors:
            diff = self.position - neighbor.position
            distance = np.linalg.norm(diff)
            if distance > 0:
                separation += diff / (distance * distance)  # 距离越近，力越大
        return separation

    def calculate_alignment(self):
        """计算对齐力（与邻居速度方向一致）"""
        if not self.neighbors:
            return np.zeros(2, dtype=float)  # 确保是浮点数

        avg_velocity = np.zeros(2, dtype=float)  # 确保是浮点数
        for neighbor in self.neighbors:
            avg_velocity += neighbor.velocity
        avg_velocity /= len(self.neighbors)

        return (avg_velocity - self.velocity) * 0.1  # 对齐系数

    def calculate_cohesion(self):
        """计算凝聚力（向邻居中心移动）"""
        if not self.neighbors:
            return np.zeros(2, dtype=float)  # 确保是浮点数

        center = np.zeros(2, dtype=float)  # 确保是浮点数
        for neighbor in self.neighbors:
            center += neighbor.position
        center /= len(self.neighbors)

        return (center - self.position) * 0.05  # 凝聚系数

    def scan_area(self, dt):
        """扫描区域行为"""
        self.velocity = np.array([0.0, 0.0], dtype=float)  # 停止移动，确保是浮点数
        self.scan_time += dt

        # 模拟发现目标（随机）
        if np.random.random() < 0.1:  # 10% 的概率发现目标
            target_pos = self.position + np.random.uniform(-20, 20, 2)
            # 确保目标在合理范围内
            target_pos = np.clip(target_pos, 0, 100)
            self.targets.append(target_pos)

        # 扫描完成后继续巡逻
        if self.scan_time > self.scan_duration:
            self.state = "patrolling"

    def avoid_obstacles(self, dt):
        """避障行为"""
        if not hasattr(self, 'avoid_target'):
            self.state = "patrolling"
            return

        # 计算远离障碍物的方向
        obstacle_dir = self.position - self.avoid_target.position
        distance = np.linalg.norm(obstacle_dir)

        if distance > 0:
            obstacle_dir = obstacle_dir / distance

        # 计算切线方向（绕过障碍物）
        tangent = np.array([-obstacle_dir[1], obstacle_dir[0]], dtype=float)  # 确保是浮点数

        # 组合方向（主要远离，部分切线）
        desired_direction = obstacle_dir * 2.0 + tangent * 1.0
        desired_direction = desired_direction / np.linalg.norm(desired_direction)

        # 设置速度
        self.velocity = desired_direction * self.max_speed * 0.8  # 避障时稍慢
        self.position += self.velocity * dt

        # 检查是否已经避开障碍物
        if distance > self.avoid_target.radius + 20:
            self.state = "patrolling"
            delattr(self, 'avoid_target')

    def return_to_base(self, dt):
        """返回基地行为"""
        # 返回第一个航点（基地）
        target = self.waypoints[0]
        direction = target - self.position
        distance = np.linalg.norm(direction)

        if distance < 2.0:  # 到达基地
            # 充电
            self.battery_level = min(100, self.battery_level + 20 * dt)
            if self.battery_level >= 80:
                self.state = "patrolling"
            return

        if distance > 0:
            direction = direction / distance

        # 计算期望速度
        desired_velocity = direction * self.max_speed

        # 应用分离行为（避免与其他无人机碰撞）
        if self.neighbors:
            separation = self.calculate_separation()
            desired_velocity += separation * 1.5

            # 限制速度
            speed = np.linalg.norm(desired_velocity)
            if speed > self.max_speed:
                desired_velocity = desired_velocity / speed * self.max_speed

        # 计算转向力
        steering = desired_velocity - self.velocity
        steering_norm = np.linalg.norm(steering)

        if steering_norm > self.max_acceleration:
            steering = steering / steering_norm * self.max_acceleration

        # 更新速度和位置
        self.velocity += steering * dt
        self.position += self.velocity * dt


class Obstacle:
    def __init__(self, position, radius):
        self.position = np.array(position, dtype=float)  # 确保是浮点数
        self.radius = radius


class SwarmController:
    def __init__(self, area_size=(100, 100)):
        self.area_size = area_size
        self.drones = []
        self.obstacles = []
        self.targets = []  # 所有无人机发现的目标
        self.time = 0
        self.dt = 0.1  # 时间步长

    def add_drone(self, drone):
        """添加无人机到集群"""
        self.drones.append(drone)

    def add_obstacle(self, obstacle):
        """添加障碍物"""
        self.obstacles.append(obstacle)

    def update(self):
        """更新所有无人机状态"""
        for drone in self.drones:
            drone.update(self.drones, self.obstacles, self.dt)

        # 收集所有目标
        self.targets = []
        for drone in self.drones:
            self.targets.extend(drone.targets)

        self.time += self.dt

    def simulate(self, duration):
        """模拟一段时间"""
        steps = int(duration / self.dt)
        for _ in range(steps):
            self.update()

    def get_positions(self):
        """获取所有无人机的位置"""
        return [drone.position for drone in self.drones]

    def get_paths(self):
        """获取所有无人机的路径"""
        return [drone.path for drone in self.drones]


# 创建模拟环境
def create_simulation():
    # 创建集群控制器
    controller = SwarmController(area_size=(100, 100))

    # 创建两架无人机
    drone1 = Drone(1, [10.0, 10.0])  # 确保使用浮点数
    drone2 = Drone(2, [90.0, 10.0])  # 确保使用浮点数

    # 设置巡逻路径点 (矩形巡逻路径)
    waypoints = [
        np.array([10.0, 10.0], dtype=float),  # 确保使用浮点数
        np.array([90.0, 10.0], dtype=float),  # 确保使用浮点数
        np.array([90.0, 90.0], dtype=float),  # 确保使用浮点数
        np.array([10.0, 90.0], dtype=float)  # 确保使用浮点数
    ]

    drone1.set_waypoints(waypoints)
    drone2.set_waypoints(waypoints[::-1])  # 反向路径，实现交叉巡逻

    # 添加无人机到集群
    controller.add_drone(drone1)
    controller.add_drone(drone2)

    # 添加障碍物
    obstacles = [
        Obstacle([30.0, 40.0], 8),  # 确保使用浮点数
        Obstacle([60.0, 30.0], 10),  # 确保使用浮点数
        Obstacle([50.0, 70.0], 12),  # 确保使用浮点数
        Obstacle([80.0, 60.0], 7)  # 确保使用浮点数
    ]

    for obs in obstacles:
        controller.add_obstacle(obs)

    return controller


# 简化可视化函数
def visualize_simulation(controller, duration):
    # 使用非交互式后端
    plt.switch_backend('Agg')  # 使用非交互式后端

    fig, ax = plt.subplots(figsize=(12, 10))
    ax.set_xlim(0, controller.area_size[0])
    ax.set_ylim(0, controller.area_size[1])
    ax.set_title('Enhanced Two-Drone Cooperative Reconnaissance Patrol')
    ax.set_xlabel('X Position')
    ax.set_ylabel('Y Position')

    # 绘制路径点
    waypoints = controller.drones[0].waypoints
    for i, point in enumerate(waypoints):
        ax.plot(point[0], point[1], 'ro', markersize=8)
        ax.text(point[0] + 2, point[1] + 2, f'WP{i + 1}', fontsize=12)

    # 绘制障碍物
    for obs in controller.obstacles:
        circle = Circle(obs.position, obs.radius, color='gray', alpha=0.7)
        ax.add_patch(circle)

    # 初始化无人机位置和路径
    drone_positions = controller.get_positions()
    drone_paths = controller.get_paths()

    drones_scat = ax.scatter([p[0] for p in drone_positions],
                             [p[1] for p in drone_positions],
                             c=['red', 'blue'], s=100, label='Drones')

    path_lines = []
    for i, path in enumerate(drone_paths):
        line, = ax.plot([p[0] for p in path], [p[1] for p in path],
                        linestyle='-', linewidth=1, alpha=0.5,
                        color='red' if i == 0 else 'blue')
        path_lines.append(line)

    # 绘制传感器范围
    sensor_circles = []
    for i, drone in enumerate(controller.drones):
        circle = Circle(drone_positions[i], drone.sensor_range,
                        fill=False, linestyle=':',
                        color='red' if i == 0 else 'blue', alpha=0.5)
        ax.add_patch(circle)
        sensor_circles.append(circle)

    # 绘制通信范围
    comm_circles = []
    for i, drone in enumerate(controller.drones):
        circle = Circle(drone_positions[i], drone.communication_range,
                        fill=False, linestyle='--',
                        color='red' if i == 0 else 'blue', alpha=0.3)
        ax.add_patch(circle)
        comm_circles.append(circle)

    # 绘制目标
    targets_scat = ax.scatter([], [], c='green', marker='*', s=100, label='Targets')

    # 添加图例
    ax.legend()

    # 信息文本框
    info_text = ax.text(2, 95, '', fontsize=10, bbox=dict(facecolor='white', alpha=0.7))

    # 电池状态条
    battery_bars = []
    for i, drone in enumerate(controller.drones):
        bar = Rectangle((5 + i * 45, 97), 40, 2,
                        color='green' if drone.battery_level > 50 else
                        'orange' if drone.battery_level > 20 else 'red')
        ax.add_patch(bar)
        battery_bars.append(bar)
        ax.text(5 + i * 45, 100, f'Drone {i + 1} Battery', fontsize=8)

    def update(frame):
        # 更新模拟
        controller.update()

        # 更新无人机位置
        positions = controller.get_positions()
        drones_scat.set_offsets(positions)

        # 更新路径
        paths = controller.get_paths()
        for i, line in enumerate(path_lines):
            if len(paths[i]) > 0:
                line.set_data([p[0] for p in paths[i]], [p[1] for p in paths[i]])

        # 更新传感器和通信范围圆
        for i, circle in enumerate(sensor_circles):
            circle.center = positions[i]
        for i, circle in enumerate(comm_circles):
            circle.center = positions[i]

        # 更新目标
        if controller.targets:
            targets_scat.set_offsets(controller.targets)

        # 更新信息文本
        info = f"Time: {controller.time:.1f}s\n"
        for i, drone in enumerate(controller.drones):
            info += f"Drone {i + 1}: {drone.state} "
            info += f"Battery: {drone.battery_level:.1f}%\n"
            info += f"Neighbors: {len(drone.neighbors)} "
            info += f"Targets: {len(drone.targets)}\n"
        info_text.set_text(info)

        # 更新电池状态条
        for i, bar in enumerate(battery_bars):
            bar.set_width(controller.drones[i].battery_level * 0.4)  # 缩放以适应显示
            bar.set_color('green' if controller.drones[i].battery_level > 50 else
                          'orange' if controller.drones[i].battery_level > 20 else 'red')

        return drones_scat, *path_lines, *sensor_circles, *comm_circles, targets_scat, info_text, *battery_bars

    # 创建动画
    ani = FuncAnimation(fig, update, frames=int(duration / controller.dt),
                        interval=100, blit=False, repeat=False)  # 增加间隔，减少帧率

    # 保存动画为GIF文件而不是MP4
    print("Saving animation as GIF...")
    try:
        ani.save('drone_simulation.gif', writer='pillow', fps=10, dpi=100)
        print("Animation saved as drone_simulation.gif")
    except Exception as e:
        print(f"Error saving animation: {e}")
        print("Trying to save as PNG frames...")
        # 如果GIF保存失败，尝试保存为PNG帧
        for i in range(min(100, int(duration / controller.dt))):  # 限制帧数
            if i % 10 == 0:  # 每10帧保存一次
                update(i)
                plt.savefig(f'frame_{i:03d}.png')
        print("Frames saved as PNG files")

    plt.close()


# 替代方案：不使用动画，而是逐步显示
def step_by_step_simulation(controller, duration):
    """逐步模拟并显示结果"""
    # 使用交互式后端
    plt.switch_backend('TkAgg')  # 使用交互式后端

    steps = int(duration / controller.dt)
    plot_interval = 20  # 每20步绘制一次

    for step in range(steps):
        controller.update()

        # 定期绘制
        if step % plot_interval == 0:
            plt.clf()
            plt.xlim(0, controller.area_size[0])
            plt.ylim(0, controller.area_size[1])
            plt.title(f'Two-Drone Cooperative Patrol - Time: {controller.time:.1f}s')
            plt.xlabel('X Position')
            plt.ylabel('Y Position')

            # 绘制路径点
            waypoints = controller.drones[0].waypoints
            for i, point in enumerate(waypoints):
                plt.plot(point[0], point[1], 'ro', markersize=8)
                plt.text(point[0] + 2, point[1] + 2, f'WP{i + 1}', fontsize=12)

            # 绘制障碍物
            for obs in controller.obstacles:
                circle = plt.Circle(obs.position, obs.radius, color='gray', alpha=0.7)
                plt.gca().add_patch(circle)

            # 绘制无人机路径
            colors = ['red', 'blue']
            for i, drone in enumerate(controller.drones):
                if drone.path:
                    path = np.array(drone.path)
                    plt.plot(path[:, 0], path[:, 1], color=colors[i], linestyle='-', linewidth=1, alpha=0.5)

                # 绘制无人机当前位置
                plt.plot(drone.position[0], drone.position[1], 'o', color=colors[i], markersize=10)

                # 绘制传感器范围
                sensor_circle = plt.Circle(drone.position, drone.sensor_range,
                                           fill=False, linestyle=':',
                                           color=colors[i], alpha=0.5)
                plt.gca().add_patch(sensor_circle)

                # 绘制通信范围
                comm_circle = plt.Circle(drone.position, drone.communication_range,
                                         fill=False, linestyle='--',
                                         color=colors[i], alpha=0.3)
                plt.gca().add_patch(comm_circle)

            # 绘制目标
            if controller.targets:
                targets = np.array(controller.targets)
                plt.scatter(targets[:, 0], targets[:, 1], c='green', marker='*', s=100)

            plt.grid(True)
            plt.pause(0.01)  # 短暂暂停以更新图形

    plt.show()


# 运行模拟
if __name__ == "__main__":
    print("Starting enhanced two-drone cooperative reconnaissance patrol simulation...")
    controller = create_simulation()

    # 选择一种可视化方式：

    # 1. 使用简化动画（保存为GIF文件）
    visualize_simulation(controller, duration=30)  # 30秒动画

    # 2. 或者使用逐步模拟（显示静态图像序列）
    # step_by_step_simulation(controller, duration=60)  # 60秒模拟

    print("Simulation completed.")